// Copyright 2017 the V8 project authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef V8_COMPILER_PERSISTENT_MAP_H_
#define V8_COMPILER_PERSISTENT_MAP_H_

#include <array>
#include <tuple>

#include "src/base/functional.h"
#include "src/zone/zone-containers.h"

namespace v8 {
namespace internal {
    namespace compiler {

        // PersistentMap is a persistent map datastructure based on hash trees (a binary
        // tree using the bits of a hash value as addresses). The map is a conceptually
        // infinite: All keys are initially mapped to a default value, values are
        // deleted by overwriting them with the default value. The iterators produce
        // exactly the keys that are not the default value. The hash values should have
        // high variance in their high bits, so dense integers are a bad choice.
        // Complexity:
        // - Copy and assignment: O(1)
        // - access: O(log n)
        // - update: O(log n) time and space
        // - iteration: amortized O(1) per step
        // - Zip: O(n)
        // - equality check: O(n)
        // TODO(tebbi): Cache map transitions to avoid re-allocation of the same map.
        // TODO(tebbi): Implement an O(1) equality check based on hash consing or
        //              something similar.
        template <class Key, class Value, class Hasher = base::hash<Key>>
        class PersistentMap {
        public:
            using key_type = Key;
            using mapped_type = Value;
            using value_type = std::pair<Key, Value>;

        private:
            static constexpr size_t kHashBits = 32;
            enum Bit : int { kLeft = 0,
                kRight = 1 };

            // Access hash bits starting from the high bits and compare them according to
            // their unsigned value. This way, the order in the hash tree is compatible
            // with numeric hash comparisons.
            class HashValue;

            struct KeyValue : std::pair<Key, Value> {
                const Key& key() const { return this->first; }
                const Value& value() const { return this->second; }
                using std::pair<Key, Value>::pair;
            };

            struct FocusedTree;

        public:
            // Depth of the last added element. This is a cheap estimate for the size of
            // the hash tree.
            size_t last_depth() const
            {
                if (tree_) {
                    return tree_->length;
                } else {
                    return 0;
                }
            }

            const Value& Get(const Key& key) const
            {
                HashValue key_hash = HashValue(Hasher()(key));
                const FocusedTree* tree = FindHash(key_hash);
                return GetFocusedValue(tree, key);
            }

            // Add or overwrite an existing key-value pair.
            void Set(Key key, Value value);

            bool operator==(const PersistentMap& other) const
            {
                if (tree_ == other.tree_)
                    return true;
                if (def_value_ != other.def_value_)
                    return false;
                for (const std::tuple<Key, Value, Value>& triple : Zip(other)) {
                    if (std::get<1>(triple) != std::get<2>(triple))
                        return false;
                }
                return true;
            }

            bool operator!=(const PersistentMap& other) const
            {
                return !(*this == other);
            }

            // The iterator produces key-value pairs in the lexicographical order of
            // hash value and key. It produces exactly the key-value pairs where the value
            // is not the default value.
            class iterator;

            iterator begin() const
            {
                if (!tree_)
                    return end();
                return iterator::begin(tree_, def_value_);
            }
            iterator end() const { return iterator::end(def_value_); }

            // Iterator to traverse two maps in lockstep, producing matching value pairs
            // for each key where at least one value is different from the respective
            // default.
            class double_iterator;

            // An iterable to iterate over the two maps in lockstep.
            struct ZipIterable {
                PersistentMap a;
                PersistentMap b;
                double_iterator begin() { return double_iterator(a.begin(), b.begin()); }
                double_iterator end() { return double_iterator(a.end(), b.end()); }
            };

            ZipIterable Zip(const PersistentMap& other) const { return { *this, other }; }

            explicit PersistentMap(Zone* zone, Value def_value = Value())
                : PersistentMap(nullptr, zone, def_value)
            {
            }

        private:
            // Find the {FocusedTree} that contains a key-value pair with key hash {hash}.
            const FocusedTree* FindHash(HashValue hash) const;

            // Find the {FocusedTree} that contains a key-value pair with key hash {hash}.
            // Output the path to this {FocusedTree} and its length. If no such
            // {FocusedTree} exists, return {nullptr} and output the path to the last node
            // with a matching hash prefix. Note that {length} is the length of the found
            // path and may be less than the length of the found {FocusedTree}.
            const FocusedTree* FindHash(HashValue hash,
                std::array<const FocusedTree*, kHashBits>* path,
                int* length) const;

            // Load value from the leaf node on the focused path of {tree}.
            const Value& GetFocusedValue(const FocusedTree* tree, const Key& key) const;

            // Return the {FocusedTree} representing the left (bit==kLeft) or right
            // (bit==kRight) child of the node on the path of {tree} at tree level
            // {level}.
            static const FocusedTree* GetChild(const FocusedTree* tree, int level,
                Bit bit);

            // Find the leftmost path in the tree, starting at the node at tree level
            // {level} on the path of {start}. Output the level of the leaf to {level} and
            // the path to {path}.
            static const FocusedTree* FindLeftmost(
                const FocusedTree* start, int* level,
                std::array<const FocusedTree*, kHashBits>* path);

            PersistentMap(const FocusedTree* tree, Zone* zone, Value def_value)
                : tree_(tree)
                , def_value_(def_value)
                , zone_(zone)
            {
            }

            const FocusedTree* tree_;
            Value def_value_;
            Zone* zone_;
        };

        // This structure represents a hash tree with one focused path to a specific
        // leaf. For the focused leaf, it stores key, value and key hash. The path is
        // defined by the hash bits of the focused leaf. In a traditional tree
        // datastructure, the nodes of a path form a linked list with the values being
        // the pointers outside of this path. Instead of storing all of these nodes,
        // we store an array of the pointers pointing outside of the path. This is
        // similar to the stack used when doing DFS traversal of a tree. The hash of
        // the leaf is used to know if the pointers point to the left or the
        // right of the path. As there is no explicit representation of a tree node,
        // this structure also represents all the nodes on its path. The intended node
        // depends on the tree depth, which is always clear from the referencing
        // context. So the pointer to a {FocusedTree} stored in the
        // {PersistentMap.tree_} always references the root, while a pointer from a
        // focused node of another {FocusedTree} always references to one tree level
        // lower than before.
        template <class Key, class Value, class Hasher>
        struct PersistentMap<Key, Value, Hasher>::FocusedTree {
            KeyValue key_value;
            // The depth of the focused path, that is, the number of pointers stored in
            // this structure.
            int8_t length;
            HashValue key_hash;
            // Out-of-line storage for hash collisions.
            const ZoneMap<Key, Value>* more;
            using more_iterator = typename ZoneMap<Key, Value>::const_iterator;
            // {path_array} has to be the last member: To store an array inline, we
            // over-allocate memory for this structure and access memory beyond
            // {path_array}. This corresponds to a flexible array member as defined in
            // C99.
            const FocusedTree* path_array[1];
            const FocusedTree*& path(int i)
            {
                DCHECK(i < length);
                return reinterpret_cast<const FocusedTree**>(
                    reinterpret_cast<byte*>(this) + offsetof(FocusedTree, path_array))[i];
            }
            const FocusedTree* path(int i) const
            {
                DCHECK(i < length);
                return reinterpret_cast<const FocusedTree* const*>(
                    reinterpret_cast<const byte*>(this) + offsetof(FocusedTree, path_array))[i];
            }
        };

        template <class Key, class Value, class Hasher>
        class PersistentMap<Key, Value, Hasher>::HashValue {
        public:
            explicit HashValue(size_t hash)
                : bits_(static_cast<uint32_t>(hash))
            {
            }

            Bit operator[](int pos) const
            {
                DCHECK_LT(pos, kHashBits);
                return bits_ & (static_cast<decltype(bits_)>(1) << (kHashBits - pos - 1))
                    ? kRight
                    : kLeft;
            }

            bool operator<(HashValue other) const { return bits_ < other.bits_; }
            bool operator==(HashValue other) const { return bits_ == other.bits_; }
            bool operator!=(HashValue other) const { return bits_ != other.bits_; }
            HashValue operator^(HashValue other) const
            {
                return HashValue(bits_ ^ other.bits_);
            }

        private:
            static_assert(sizeof(uint32_t) * 8 == kHashBits, "wrong type for bits_");
            uint32_t bits_;
        };

        template <class Key, class Value, class Hasher>
        class PersistentMap<Key, Value, Hasher>::iterator {
        public:
            const value_type operator*() const
            {
                if (current_->more) {
                    return *more_iter_;
                } else {
                    return current_->key_value;
                }
            }

            iterator& operator++()
            {
                do {
                    if (!current_) {
                        // Iterator is past the end.
                        return *this;
                    }
                    if (current_->more) {
                        DCHECK(more_iter_ != current_->more->end());
                        ++more_iter_;
                        if (more_iter_ != current_->more->end())
                            return *this;
                    }
                    if (level_ == 0) {
                        *this = end(def_value_);
                        return *this;
                    }
                    --level_;
                    while (current_->key_hash[level_] == kRight || path_[level_] == nullptr) {
                        if (level_ == 0) {
                            *this = end(def_value_);
                            return *this;
                        }
                        --level_;
                    }
                    const FocusedTree* first_right_alternative = path_[level_];
                    level_++;
                    current_ = FindLeftmost(first_right_alternative, &level_, &path_);
                    if (current_->more) {
                        more_iter_ = current_->more->begin();
                    }
                } while (!((**this).second != def_value()));
                return *this;
            }

            bool operator==(const iterator& other) const
            {
                if (is_end())
                    return other.is_end();
                if (other.is_end())
                    return false;
                if (current_->key_hash != other.current_->key_hash) {
                    return false;
                } else {
                    return (**this).first == (*other).first;
                }
            }
            bool operator!=(const iterator& other) const { return !(*this == other); }

            bool operator<(const iterator& other) const
            {
                if (is_end())
                    return false;
                if (other.is_end())
                    return true;
                if (current_->key_hash == other.current_->key_hash) {
                    return (**this).first < (*other).first;
                } else {
                    return current_->key_hash < other.current_->key_hash;
                }
            }

            bool is_end() const { return current_ == nullptr; }

            const Value& def_value() { return def_value_; }

            static iterator begin(const FocusedTree* tree, Value def_value)
            {
                iterator i(def_value);
                i.current_ = FindLeftmost(tree, &i.level_, &i.path_);
                if (i.current_->more) {
                    i.more_iter_ = i.current_->more->begin();
                }
                // Skip entries with default value. PersistentMap iterators must never point
                // to a default value.
                while (!i.is_end() && !((*i).second != def_value))
                    ++i;
                return i;
            }

            static iterator end(Value def_value) { return iterator(def_value); }

        private:
            int level_;
            typename FocusedTree::more_iterator more_iter_;
            const FocusedTree* current_;
            std::array<const FocusedTree*, kHashBits> path_;
            Value def_value_;

            explicit iterator(Value def_value)
                : level_(0)
                , current_(nullptr)
                , def_value_(def_value)
            {
            }
        };

        template <class Key, class Value, class Hasher>
        class PersistentMap<Key, Value, Hasher>::double_iterator {
        public:
            std::tuple<Key, Value, Value> operator*()
            {
                if (first_current_) {
                    auto pair = *first_;
                    return std::make_tuple(
                        pair.first, pair.second,
                        second_current_ ? (*second_).second : second_.def_value());
                } else {
                    DCHECK(second_current_);
                    auto pair = *second_;
                    return std::make_tuple(pair.first, first_.def_value(), pair.second);
                }
            }

            double_iterator& operator++()
            {
#ifdef DEBUG
                iterator old_first = first_;
                iterator old_second = second_;
#endif
                if (first_current_) {
                    ++first_;
                    DCHECK(old_first < first_);
                }
                if (second_current_) {
                    ++second_;
                    DCHECK(old_second < second_);
                }
                return *this = double_iterator(first_, second_);
            }

            double_iterator(iterator first, iterator second)
                : first_(first)
                , second_(second)
            {
                if (first_ == second_) {
                    first_current_ = second_current_ = true;
                } else if (first_ < second_) {
                    first_current_ = true;
                    second_current_ = false;
                } else {
                    DCHECK(second_ < first_);
                    first_current_ = false;
                    second_current_ = true;
                }
            }

            bool operator!=(const double_iterator& other)
            {
                return first_ != other.first_ || second_ != other.second_;
            }

            bool is_end() const { return first_.is_end() && second_.is_end(); }

        private:
            iterator first_;
            iterator second_;
            bool first_current_;
            bool second_current_;
        };

        template <class Key, class Value, class Hasher>
        void PersistentMap<Key, Value, Hasher>::Set(Key key, Value value)
        {
            HashValue key_hash = HashValue(Hasher()(key));
            std::array<const FocusedTree*, kHashBits> path;
            int length = 0;
            const FocusedTree* old = FindHash(key_hash, &path, &length);
            ZoneMap<Key, Value>* more = nullptr;
            if (!(GetFocusedValue(old, key) != value))
                return;
            if (old && !(old->more == nullptr && old->key_value.key() == key)) {
                more = new (zone_->New(sizeof(*more))) ZoneMap<Key, Value>(zone_);
                if (old->more) {
                    *more = *old->more;
                } else {
                    (*more)[old->key_value.key()] = old->key_value.value();
                }
                (*more)[key] = value;
            }
            FocusedTree* tree = new (zone_->New(sizeof(FocusedTree) + std::max(0, length - 1) * sizeof(const FocusedTree*)))
                FocusedTree { KeyValue(std::move(key), std::move(value)),
                    static_cast<int8_t>(length),
                    key_hash,
                    more,
                    {} };
            for (int i = 0; i < length; ++i) {
                tree->path(i) = path[i];
            }
            *this = PersistentMap(tree, zone_, def_value_);
        }

        template <class Key, class Value, class Hasher>
        const typename PersistentMap<Key, Value, Hasher>::FocusedTree*
        PersistentMap<Key, Value, Hasher>::FindHash(HashValue hash) const
        {
            const FocusedTree* tree = tree_;
            int level = 0;
            while (tree && hash != tree->key_hash) {
                while ((hash ^ tree->key_hash)[level] == 0) {
                    ++level;
                }
                tree = level < tree->length ? tree->path(level) : nullptr;
                ++level;
            }
            return tree;
        }

        template <class Key, class Value, class Hasher>
        const typename PersistentMap<Key, Value, Hasher>::FocusedTree*
        PersistentMap<Key, Value, Hasher>::FindHash(
            HashValue hash, std::array<const FocusedTree*, kHashBits>* path,
            int* length) const
        {
            const FocusedTree* tree = tree_;
            int level = 0;
            while (tree && hash != tree->key_hash) {
                int map_length = tree->length;
                while ((hash ^ tree->key_hash)[level] == 0) {
                    (*path)[level] = level < map_length ? tree->path(level) : nullptr;
                    ++level;
                }
                (*path)[level] = tree;
                tree = level < tree->length ? tree->path(level) : nullptr;
                ++level;
            }
            if (tree) {
                while (level < tree->length) {
                    (*path)[level] = tree->path(level);
                    ++level;
                }
            }
            *length = level;
            return tree;
        }

        template <class Key, class Value, class Hasher>
        const Value& PersistentMap<Key, Value, Hasher>::GetFocusedValue(
            const FocusedTree* tree, const Key& key) const
        {
            if (!tree) {
                return def_value_;
            }
            if (tree->more) {
                auto it = tree->more->find(key);
                if (it == tree->more->end())
                    return def_value_;
                else
                    return it->second;
            } else {
                if (key == tree->key_value.key()) {
                    return tree->key_value.value();
                } else {
                    return def_value_;
                }
            }
        }

        template <class Key, class Value, class Hasher>
        const typename PersistentMap<Key, Value, Hasher>::FocusedTree*
        PersistentMap<Key, Value, Hasher>::GetChild(const FocusedTree* tree, int level,
            Bit bit)
        {
            if (tree->key_hash[level] == bit) {
                return tree;
            } else if (level < tree->length) {
                return tree->path(level);
            } else {
                return nullptr;
            }
        }

        template <class Key, class Value, class Hasher>
        const typename PersistentMap<Key, Value, Hasher>::FocusedTree*
        PersistentMap<Key, Value, Hasher>::FindLeftmost(
            const FocusedTree* start, int* level,
            std::array<const FocusedTree*, kHashBits>* path)
        {
            const FocusedTree* current = start;
            while (*level < current->length) {
                if (const FocusedTree* child = GetChild(current, *level, kLeft)) {
                    (*path)[*level] = GetChild(current, *level, kRight);
                    current = child;
                    ++*level;
                } else if (const FocusedTree* child = GetChild(current, *level, kRight)) {
                    (*path)[*level] = GetChild(current, *level, kLeft);
                    current = child;
                    ++*level;
                } else {
                    UNREACHABLE();
                }
            }
            return current;
        }

        template <class Key, class Value, class Hasher>
        std::ostream& operator<<(std::ostream& os,
            const PersistentMap<Key, Value, Hasher>& map)
        {
            os << "{";
            bool first = true;
            for (auto pair : map) {
                if (!first)
                    os << ", ";
                first = false;
                os << pair.first << ": " << pair.second;
            }
            return os << "}";
        }

    } // namespace compiler
} // namespace internal
} // namespace v8

#endif // V8_COMPILER_PERSISTENT_MAP_H_
